# ---------------------------------------------#
#   运行前一定要修改classes
#   如果生成的2007_train.txt里面没有目标信息
#   那么就是因为classes没有设定正确
# ---------------------------------------------#
import xml.etree.ElementTree as ET
from os import getcwd

sets = [('2007', 'train'), ('2007', 'val'), ('2007', 'test')]

classes = ["aeroplane", "bicycle", "bird", "boat", "bottle", "bus", "car", "cat", "chair", "cow", "diningtable", "dog",
           "horse", "motorbike", "person", "pottedplant", "sheep", "sofa", "train", "tvmonitor"]


def convert_annotation(year, image_id, list_file):
    in_file = open('VOCdevkit/VOC%s/Annotations/%s.xml' % (year, image_id), encoding='utf-8')
    tree = ET.parse(in_file)  # 载入当前xml文件的数据
    root = tree.getroot()  # 获取其根节点，tree.getroot()返回的是一个obj地址

    """
    一个object下的内容
    <object>
		<name>dog</name>
		<pose>Left</pose>
		<truncated>1</truncated>
		<difficult>0</difficult>
		<bndbox>
			<xmin>48</xmin>
			<ymin>240</ymin>
			<xmax>195</xmax>
			<ymax>371</ymax>
		</bndbox>
	</object>
    """
    for obj in root.iter('object'):  # 递归查找所有子节点
        difficult = 0
        if obj.find('difficult') is not None:  # obj的子节点difficult节点的值不为空
            difficult = obj.find('difficult').text  # obj的子节点difficult节点的值

        cls = obj.find('name').text  # 获得当前object的类别名
        if cls not in classes or int(difficult) == 1:  # 如果该类别不属于事先定好的类别或者
            continue
        cls_id = classes.index(cls)  # 检测当中是否包含字符串cls
        xmlbox = obj.find('bndbox')  # 找到边界框内容
        b = (int(xmlbox.find('xmin').text), int(xmlbox.find('ymin').text), int(xmlbox.find('xmax').text),
             int(xmlbox.find('ymax').text))  # 获取边界框的各个值
        list_file.write(" " + ",".join([str(a) for a in b]) + ',' + str(cls_id))  # 写入文件


# 返回当前工作目录
wd = getcwd()

for year, image_set in sets:
    # read()的意思是从文件指针所在的位置，读到文件结尾。括号内添加数字代表读取几个字节
    # strip()默认删除字符串头和尾的空白字符
    # 所以第一行的赋值意思是将当前目录下的文件内容按空白字符分割并存储为列表
    image_ids = open('VOCdevkit/VOC%s/ImageSets/Main/%s.txt' % (year, image_set)).read().strip().split()

    # 创建并写入文件list_file。2007_test.txt和2007_val.txt无内容是正常的
    list_file = open('%s_%s.txt' % (year, image_set), 'w')
    for image_id in image_ids:
        print('%s/VOCdevkit/VOC%s/JPEGImages/%s.jpg' % (wd, year, image_id))
        list_file.write('%s/VOCdevkit/VOC%s/JPEGImages/%s.jpg' % (wd, year, image_id))
        convert_annotation(year, image_id, list_file)
        list_file.write('\n')
    list_file.close()
